import random

N = 8
k = 3


def get_conflicts_num(status):
    num = []
    for q in range(k):
        n = 0
        for i in range(N):
            for j in range(i + 1, N):
                if status[q][i] == status[q][j]:
                    n += 1
                if abs(status[q][i] - status[q][j]) == (j - i):
                    n += 1
        num.append(n)
    return num


def get_conflicts_single(status):
    num = 0
    for i in range(N):
        for j in range(i + 1, N):
            if status[i] == status[j]:
                num += 1
            if abs(status[i] - status[j]) == (j - i):
                num += 1
    return num


def init_status():
    status = []
    temp = []
    for i in range(k):
        for j in range(N):
            temp.append(random.randint(0, N - 1))
        status.append(list(temp))
        temp.clear()
    return status


def beam_search(status):
    dic = {}
    for x in range(k):
        for y in range(N):
            for z in range(N):
                if status[x][y] == z:
                    continue
                new_status = list(status[x])
                new_status[y] = z
                dic[(x, y, z)] = get_conflicts_single(new_status)  # 生成所有后继
    ans = []
    # 局部束搜索
    # for i in range(k):  # 找k个最优值作为后继
    #     conflicts = float('inf')
    #     for a, b in dic.items():
    #         if b < conflicts:
    #             conflicts = b
    #     for a, b in dic.items():
    #         if b == conflicts:
    #             ans.append(a)
    #             dic.pop(a)
    #             break
    # if len(ans) > 0:
    #     return_status = []
    #     for x in range(k):
    #         p = ans[x][0]
    #         q = ans[x][1]
    #         m = ans[x][2]
    #         temp_status = list(status[p])
    #         temp_status[q] = m
    #         return_status.append(temp_status)
    #     return return_status
    # endof 局部束搜索

    # 随机束搜索
    while len(ans) < k:
        conflicts = float('inf')
        for a, b in dic.items():
            if b < conflicts:
                conflicts = b
        for a, b in dic.items():
            if b == conflicts:
                ans.append(a)
    return_status = []
    for i in range(k):
        rdnum = random.randint(0, len(ans)-1)
        p = ans[rdnum][0]
        q = ans[rdnum][1]
        m = ans[rdnum][2]
        temp_status = list(status[p])
        temp_status[q] = m
        return_status.append(temp_status)
    return return_status
    # endof 随机束搜索


def solve():
    status = init_status()
    print("initial status:")
    for i in range(k):
        print(status[i])
    conflicts = get_conflicts_num(status)
    while conflicts.count(0) == 0:
        status = beam_search(status)
        conflicts = get_conflicts_num(status)
        print("status:")
        for i in range(k):
            print(status[i])
            print("conflicts:", get_conflicts_single(status[i]))
    n = conflicts.index(0)
    print("final ans:")
    print(status[n])


if __name__ == '__main__':
    solve()
